package com.chenfan.common.extension.util.wrapper;

import com.chenfan.common.extension.context.TraceContext;
import com.chenfan.common.extension.context.TraceContextHolder;
import org.slf4j.MDC;

import java.util.Map;
import java.util.function.Supplier;

/**
 * 封装Supplier接口
 * @author wulg
 * @date 2022-04-26
 **/
public class SupplierWrapper<T> implements Supplier<T> {

    private Supplier<T> supplier;

    private final Map<String, String> mdcContext = MDC.getCopyOfContextMap();

    private final TraceContext traceContext = TraceContextHolder.getCurrentContext();

    public SupplierWrapper(Supplier<T> supplier) {
        this.supplier = supplier;
    }

    @Override
    public T get() {
        try {
            if(this.mdcContext != null){
                MDC.setContextMap(this.mdcContext);
            }
            if(this.traceContext != null){
                TraceContextHolder.setContext(traceContext);
            }
            return supplier.get();
        } finally {
            if(this.mdcContext != null){
                MDC.clear();
            }
            if(this.traceContext != null){
                TraceContextHolder.clear();
            }
        }
    }
}
